# 模型训练代码
import torch  # 导入 PyTorch 库，用于深度学习模型的构建和训练
from MyData import MyDataset  # 自定义数据集类，用于加载训练数据
from torch.utils.data import DataLoader  # PyTorch 数据加载器，用于批量加载数据
from net import Model  # 自定义的模型类（基于 BERT 的分类模型）
from transformers import BertTokenizer, AdamW  # Hugging Face 提供的 BERT 分词器和优化器

# 定义设备信息（如果 GPU 可用，则使用 GPU；否则使用 CPU）
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义训练的轮次（总训练次数）
EPOCH = 30000

# 加载预训练的 BERT 分词器（指定路径为本地 BERT 模型的路径）
token = BertTokenizer.from_pretrained(r"D:\SentimentBERT\model\bert-base-chinese\models--bert-base-chinese\snapshots\c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f")

# 定义一个函数 `collate_fn`，用于对数据进行编码处理
def collate_fn(data):
    sents = [i[0] for i in data]  # 提取每个样本中的文本部分
    label = [i[1] for i in data]  # 提取每个样本中的标签部分
    # 使用分词器对文本进行编码（包括填充、截断等操作）
    data = token.batch_encode_plus(
        batch_text_or_text_pairs=sents,  # 输入的文本列表
        truncation=True,  # 是否对超出最大长度的文本进行截断
        max_length=500,  # 设置最大长度为 500
        padding="max_length",  # 将所有文本填充到最大长度
        return_tensors="pt",  # 返回 PyTorch 张量格式的数据
        return_length=True  # 返回每个文本的实际长度
    )
    input_ids = data["input_ids"]  # 获取输入的 token ID
    attention_mask = data["attention_mask"]  # 获取 attention mask（用于忽略填充部分）
    token_type_ids = data["token_type_ids"]  # 获取 token 类型 ID（用于区分句子 A 和 B）
    labels = torch.LongTensor(label)  # 将标签转换为 LongTensor 格式
    return input_ids, attention_mask, token_type_ids, labels  # 返回编码后的数据

# 创建训练数据集
train_dataset = MyDataset("train")  # 加载训练数据集（假设 `MyDataset` 已实现）

# 创建数据加载器（用于批量加载数据）
train_loader = DataLoader(
    dataset=train_dataset,  # 数据集对象
    batch_size=100,  # 每批次包含 100 条数据
    shuffle=True,  # 在每个 epoch 开始时打乱数据顺序
    drop_last=True,  # 舍弃最后一个不足 batch_size 的批次，防止形状不匹配
    collate_fn=collate_fn  # 使用自定义的 `collate_fn` 对数据进行编码
)

if __name__ == '__main__':
    # 打印设备信息（GPU 或 CPU）
    print(DEVICE)

    # 初始化模型并将模型移动到指定设备上
    model = Model().to(DEVICE)

    # 定义优化器（AdamW 是一种常用的优化算法，适合 BERT 模型）
    optimizer = AdamW(model.parameters())

    # 定义损失函数（交叉熵损失，适用于多分类任务）
    loss_func = torch.nn.CrossEntropyLoss()

    # 训练循环（外层循环：按 epoch 进行训练）
    for epoch in range(EPOCH):
        # 内层循环：按批次加载数据
        for i, (input_ids, attention_mask, token_type_ids, labels) in enumerate(train_loader):
            # 将数据移动到指定设备上
            input_ids, attention_mask, token_type_ids, labels = \
                input_ids.to(DEVICE), attention_mask.to(DEVICE), token_type_ids.to(DEVICE), labels.to(DEVICE)

            # 前向计算（将数据输入模型，得到输出）
            out = model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)

            # 根据模型输出和真实标签计算损失
            loss = loss_func(out, labels)

            # 反向传播（根据损失优化模型参数）
            optimizer.zero_grad()  # 清空梯度
            loss.backward()  # 反向传播计算梯度
            optimizer.step()  # 更新模型参数

            # 每隔 5 个批次输出一次训练信息
            if i % 5 == 0:
                out = out.argmax(dim=1)  # 获取预测结果（取输出的最大值对应的类别）
                acc = (out == labels).sum().item() / len(labels)  # 计算准确率
                print(f"epoch:{epoch}, i:{i}, loss:{loss.item()}, acc:{acc}")

        # 每训练完一轮，保存一次模型参数
        torch.save(model.state_dict(), f"params/{epoch}_bert.pth")
        print(epoch, "参数保存成功！")
#逻辑实现
#1.加载并编码训练数据
#2.初始化模型、优化器和损失函数
#3.按批次训练模型，计算损失并更新参数
#4.定期保存模型参数，便于后续推理或继续训练